#include "utreeview.h"
#include "ustyleditemdelegate.h"
#include "widgets/uscrollbar.h"
#include "uheaderview.h"

#include <QHeaderView>
#include <QKeyEvent>
#include <QPainter>
#include <QPainterPath>

constexpr int iconWidth = 15;

class UTreeViewPrivate
{
public:
    UTreeViewPrivate(UTreeView* q) : q_ptr(q) {}
    ~UTreeViewPrivate() {}

    void init();

    void _q_updateViewCheckStatus(const Qt::CheckState &status, int index);

    void _q_updateHeaderCheckStatus(const QModelIndex &index, bool isChecked);

private:
    void initUi();

private:
    Q_DECLARE_PUBLIC(UTreeView)
    UTreeView* const q_ptr = Q_NULLPTR;
    UStyledItemDelegate* delegate = Q_NULLPTR;
};

void UTreeViewPrivate::init()
{
    initUi();
}

void UTreeViewPrivate::_q_updateViewCheckStatus(const Qt::CheckState &status, int index)
{
    Q_Q(UTreeView);
    if (nullptr == q->model())
    {
        return;
    }
    for (int i = 0; i < q->model()->rowCount(); ++i)
    {
        QModelIndex modelIndex = q->model()->index(i, index);
        q->model()->setData(modelIndex, (Qt::CheckState::Checked == status));
    }
    emit q->headerCheckStatusChanged(status);
}

void UTreeViewPrivate::_q_updateHeaderCheckStatus(const QModelIndex &index, bool isChecked)
{
    Q_Q(UTreeView);

    q->updateHeaderCheckStatus();
}

void UTreeViewPrivate::initUi()
{
    Q_Q(UTreeView);

    q->setMouseTracking(true);
    q->viewport()->setMouseTracking(true);
    delegate = new UStyledItemDelegate(q);
    q->setItemDelegate(delegate);

    auto hScrollBar = new UScrollBar(q);
    q->setHorizontalScrollBar(hScrollBar);
    auto vScrollBar = new UScrollBar(q);
    q->setVerticalScrollBar(vScrollBar);
    auto header = new UHeaderView(Qt::Horizontal, q);
    q->setHeader(header);

    q->setAlternatingRowColors(true);
    q->setVerticalScrollMode(QAbstractItemView::ScrollMode::ScrollPerItem);
    q->setFrameShape(QFrame::NoFrame);
    q->setSelectionMode(QAbstractItemView::SingleSelection);
    q->setSelectionBehavior(QAbstractItemView::SelectRows);
    q->setEditTriggers(QAbstractItemView::NoEditTriggers);
    q->header()->setDefaultAlignment(Qt::AlignLeft | Qt::AlignVCenter);
    q->setTextElideMode(Qt::TextElideMode::ElideMiddle);
    q->header()->setStretchLastSection(true);
    q->header()->setSectionsMovable(false);
    q->header()->setVisible(false);
    q->setSortingEnabled(true);
    q->viewport()->setAutoFillBackground(true);

    QObject::connect(header, SIGNAL(stateChanged(const Qt::CheckState &, int)),
                     q, SLOT(_q_updateViewCheckStatus(const Qt::CheckState &, int)));

    QObject::connect(delegate, SIGNAL(checkStatusChanged(const QModelIndex &, bool)),
                     q, SLOT(_q_updateHeaderCheckStatus(const QModelIndex &, bool)));
}

UTreeView::UTreeView(QWidget* parent)
    : QTreeView(parent)
    , d_ptr(new UTreeViewPrivate(this))
{
    Q_D(UTreeView);

    d->init();
}

UTreeView::~UTreeView()
{
    Q_D(UTreeView);

    delete d;
}

void UTreeView::setItemHeight(int height)
{
    Q_D(UTreeView);

    d->delegate->setRowHeight(height);
}

void UTreeView::setCheckColumnIndex(int index)
{
    if (auto headerView = qobject_cast<UHeaderView*>(header()))
    {
        headerView->setCheckColumnIndex(index);
    }

    if (auto delegate = qobject_cast<UStyledItemDelegate*>(itemDelegate()))
    {
        delegate->setCheckColumnIndex(index);
    }
}

void UTreeView::updateHeaderCheckStatus()
{
    if (nullptr == model())
    {
        return;
    }
    auto headerView = qobject_cast<UHeaderView*>(header());
    if (nullptr == headerView)
    {
        return;
    }
    if (headerView->checkColumnIndex() < 0)
    {
        return;
    }
    QSet<bool> checkSet;
    int count = model()->rowCount();
    if (count <= 0)
    {
        headerView->setCheckState(Qt::CheckState::Unchecked);
        emit headerCheckStatusChanged(Qt::CheckState::Unchecked);
        return;
    }
    for (int i = 0; i < count; ++i)
    {
        QModelIndex modelIndex = model()->index(i, headerView->checkColumnIndex());
        checkSet.insert(modelIndex.data().toBool());
    }

    if (checkSet.count() == 1)
    {
        headerView->setCheckState(*checkSet.begin() ? Qt::CheckState::Checked : Qt::CheckState::Unchecked);
    }
    else
    {
        headerView->setCheckState(Qt::CheckState::PartiallyChecked);
    }
    emit headerCheckStatusChanged(headerView->checkState());
}

void UTreeView::keyPressEvent(QKeyEvent* event)
{
    QTreeView::keyPressEvent(event);

    if (event->key() == Qt::Key_Up || event->key() == Qt::Key_Down)
        emit clicked(currentIndex());
}

void UTreeView::drawBranches(QPainter* painter, const QRect &rect, const QModelIndex &index) const
{
    painter->save();
    painter->setRenderHint(QPainter::Antialiasing);
    painter->setOpacity(1);

    auto model = index.model();
    if (nullptr == model)
        return;

    bool hasChild = model->rowCount(index) > 0;
    bool expaned = isExpanded(index);

    QPalette palette = this->palette();
    QBrush background = palette.color(QPalette::Base);

    QPainterPath path;
    path.addRect(rect);
    painter->fillPath(path, background);

    QIcon icon;
    if (expaned)
    {
        if (hasChild)
            icon = QIcon::fromTheme("utk_expand");
    }
    else
    {
        if (hasChild)
            icon = QIcon::fromTheme("utk_collapse");
    }

    if (!icon.isNull())
    {
        QRect iconRect = QRect(rect.x() + (rect.width() - iconWidth) / 2, rect.y() + (rect.height() - iconWidth) / 2,
                               iconWidth, iconWidth);
        icon.paint(painter, iconRect, Qt::AlignCenter);
    }

    painter->restore();
}

void UTreeView::contextMenuEvent(QContextMenuEvent *event)
{
    emit contextMenu(indexAt(event->pos()), event->globalPos());
    QTreeView::contextMenuEvent(event);
}

void UTreeView::updateContent()
{
    updateGeometries();
}

#include "moc_utreeview.cpp"
